package com.gitee.wsl.coroutines.executor.sample


import com.gitee.wsl.concurrent.api.SuspendRunnableWithResult
import com.gitee.wsl.coroutines.executor.KExecutor
import com.gitee.wsl.coroutines.executor.SuspendKExecutor
import com.gitee.wsl.ext.coroutines.get
import kotlinx.atomicfu.AtomicRef
import kotlinx.atomicfu.atomic
import kotlinx.coroutines.CoroutineStart
import kotlinx.coroutines.Job
import kotlinx.coroutines.async
import kotlinx.coroutines.cancelAndJoin
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.yield


class ConflatedExecutor: SuspendKExecutor {

    private val activeTask: AtomicRef<Job?> = atomic(null)

    override suspend fun <T> executeSuspend(block: SuspendRunnableWithResult<T>): T {
        return coroutineScope {
            val newTask = async(start = CoroutineStart.LAZY) {
                block.run()
            }.also { task ->
                task.invokeOnCompletion {
                    activeTask.compareAndSet(task, null)
                }
            }

            val result: T
            while (true) {
                if (!activeTask.compareAndSet(null, newTask)) {
                    activeTask.get()?.cancelAndJoin()
                    yield()
                } else {
                    result = newTask.await()
                    break
                }
            }
            result
        }
    }

    companion object{
        val INSTANCE = ConflatedExecutor()
    }
}

fun KExecutor.Companion.conflatedExecutor() = ConflatedExecutor.INSTANCE